 #include "server/websocket_server.h"
#include "server/websocket_session.h"
#include "common/utils.h"
#include <iostream>
#include <csignal>
#include <atomic>
#include <thread>
#include <chrono>
#include <unordered_map>

std::atomic<bool> running(true);

void signal_handler(int signal) {
    std::cout << "接收到信号 " << signal << ", 正在关闭服务器..." << std::endl;
    running = false;
}

int main() {
    // 注册信号处理
    std::signal(SIGINT, signal_handler);
    std::signal(SIGTERM, signal_handler);
    
    // 创建WebSocket服务器
    websocket::WebSocketServer server(8080);
    
    // 设置连接回调
    server.setConnectionHandler([](std::shared_ptr<websocket::WebSocketSession> session) {
        utils::log_message("INFO", "客户端连接: " + session->getRemoteAddress());
        session->sendText("欢迎连接到WebSocket服务器!");
    });
    
    // 为昵称维护映射
    std::unordered_map<int, std::string> nicknames;
    // 为房间维护映射（socket -> room）
    std::unordered_map<int, std::string> rooms;

    // 设置断开连接回调
    server.setDisconnectionHandler([&nicknames, &rooms](std::shared_ptr<websocket::WebSocketSession> session) {
        utils::log_message("INFO", "客户端断开: " + session->getRemoteAddress());
        nicknames.erase(session->getSocket());
        rooms.erase(session->getSocket());
    });
    
    // 设置消息处理回调
    server.setMessageHandler([&server, &nicknames, &rooms](std::shared_ptr<websocket::WebSocketSession> session, 
                               const std::string& message) {
        utils::log_message("INFO", "收到消息: " + message);

        // 昵称设置协议：客户端发送 "NICK:xxx" 设置昵称
        if (message.rfind("NICK:", 0) == 0) {
            std::string nick = message.substr(5);
            if (nick.empty()) nick = "匿名";
            nicknames[session->getSocket()] = nick;
            session->sendText("系统: 昵称已设置为 " + nick);
            return;
        }
        
        // 房间加入：客户端发送 "JOIN:房间号"
        if (message.rfind("JOIN:", 0) == 0) {
            std::string room = message.substr(5);
            if (room.empty()) {
                session->sendText("系统: 房间号不能为空");
            } else {
                rooms[session->getSocket()] = room;
                session->sendText("系统: 已加入房间 " + room);
            }
            return;
        }

        // 处理指令
        if (message == "ping") {
            session->sendText("pong");
            return;
        } else if (message == "time") {
            session->sendText("当前时间: " + utils::get_timestamp());
            return;
        } else if (message == "stats") {
            session->sendText("服务器运行正常");
            return;
        }

        // 群聊广播：昵称优先，缺省为远端地址
        std::string sender = "匿名";
        auto it = nicknames.find(session->getSocket());
        if (it != nicknames.end() && !it->second.empty()) {
            sender = it->second;
        }
        // 按房间广播：未加入房间则提示
        auto rit = rooms.find(session->getSocket());
        if (rit == rooms.end() || rit->second.empty()) {
            session->sendText("系统: 请先加入房间");
            return;
        }
        const std::string room = rit->second;
        // 收集同房间的 sockets
        std::vector<int> targets;
        for (const auto& kv : rooms) {
            if (kv.second == room) targets.push_back(kv.first);
        }
        server.broadcastTo(targets, sender + ": " + message);
    });
    
    // 启动服务器
    if (!server.start()) {
        std::cerr << "服务器启动失败" << std::endl;
        return 1;
    }
    
    utils::log_message("INFO", "WebSocket服务器已启动，按 Ctrl+C 停止");
    
    // 主循环
    while (running) {
        std::this_thread::sleep_for(std::chrono::seconds(1));
        
        // 定期输出连接统计
        static int counter = 0;
        if (++counter % 10 == 0) {
            utils::log_message("INFO", "当前连接数: " + 
                             std::to_string(server.getConnectionCount()));
        }
    }
    
    // 停止服务器
    server.stop();
    utils::log_message("INFO", "服务器已安全关闭");
    
    return 0;
}